package tk.mybatis.web.webSocket;

import org.springframework.web.context.ContextLoader;
import org.springframework.web.context.WebApplicationContext;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;

import javax.servlet.ServletContext;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class WebSocketPool {
    private static final Map<WebSocketSession,String> userconnections = new HashMap<>();
    private static WebSocketEndPoint handler;

    public static WebSocketEndPoint getWebSocketHandler(){
        if(handler==null){
            WebApplicationContext context = ContextLoader.getCurrentWebApplicationContext() ;
            handler = context.getBean("webSocket",WebSocketEndPoint.class);
        }
        return handler;
    }

    public static void addConnection(WebSocketSession conn, String tag){
        userconnections.put(conn, tag);
    }
    public static void removeConnection(WebSocketSession conn){
        userconnections.remove(conn);
    }
    public static <T> void sendMessage(T message){
        userconnections.keySet().forEach(m-> {
            try {
                getWebSocketHandler().send(m, message);
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }
}
